import os
import nibabel as nib
import numpy as np
import matplotlib.pyplot as plt  # 新增导入

def binarize_labels(input_dir, output_dir):
    """
    将标签数据二值化处理

    参数:
    input_dir (str): 原始标签数据目录
    output_dir (str): 二值化后标签保存目录
    """
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    # 遍历输入目录中的所有文件
    for filename in sorted(os.listdir(input_dir)):
        if filename.endswith('.nii') or filename.endswith('.nii.gz'):
            # 加载标签数据
            filepath = os.path.join(input_dir, filename)
            label_nii = nib.load(filepath)
            label_array = label_nii.get_fdata()

            # 二值化处理：非零值设为1
            binary_label = (label_array > 0).astype(np.uint8)

            # 创建新的NIfTI图像并保存
            new_nii = nib.Nifti1Image(binary_label, label_nii.affine, label_nii.header)
            output_path = os.path.join(output_dir, filename)
            nib.save(new_nii, output_path)

            print(f"Processed: {filename}")

def display_slice(filepath, slice_idx=None, axis=2):
    """
    显示NIfTI文件的某个切片

    参数:
    filepath (str): NIfTI文件路径
    binarize_labels(input_label_dir, output_label_dir)

    # 示例：显示第一个文件的中间切片
    sample_file = os.path.join(output_label_dir, os.listdir(output_label_dir)[0])
    display_slice(sample_file)    slice_idx (int): 切片索引，如果为None则显示中间切片
    axis (int): 切片方向(0:矢状面, 1:冠状面, 2:横断面)
    """
    # 加载数据
    img = nib.load(filepath)
    data = img.get_fdata()

    # 如果没有指定切片索引，使用中间切片
    if slice_idx is None:
        slice_idx = data.shape[axis] // 2

    # 根据轴选择切片
    if axis == 0:
        slice_data = data[slice_idx, :, :]
    elif axis == 1:
        slice_data = data[:, slice_idx, :]
    else:
        slice_data = data[:, :, slice_idx]

    # 显示切片
    plt.figure(figsize=(8, 8))
    plt.imshow(slice_data.T, cmap='gray', origin='lower')
    plt.title(f"Slice {slice_idx} (Axis {['Sagittal', 'Coronal', 'Axial'][axis]})")
    plt.colorbar()
    plt.show()

if __name__ == "__main__":
    # # 设置输入输出路径
    # input_label_dir = r"D:\dataset\traindata\label"  # 原始标签目录
    # output_label_dir = r"D:\dataset\traindata\label_binary"  # 二值化标签保存目录
    #
    # # 执行二值化处理
    # binarize_labels(input_label_dir, output_label_dir)
    #
    # # 示例：显示第一个文件的中间切片
    # sample_file = os.path.join(input_label_dir, os.listdir(input_label_dir)[0])
    # display_slice(sample_file)

    # 显示特定文件的第50个横断面切片
    display_slice(r"E:\data\traindata\label\segmentation-3.nii", slice_idx=100, axis=0)